import matplotlib.pyplot as plt
import numpy as np
# import json

# from pathlib import Path
from multiprocessing import Queue
from queue import Empty

from tensorflow.keras.callbacks import BaseLogger
from matplotlib.animation import FuncAnimation


class Train_Monitor(BaseLogger):
    def __init__(self, que: Queue, e_num: int):
        """


        :param que: 进程队列
        :param e_num: epoch数
        """
        super().__init__()
        self.e_num = e_num
        self.record = {}
        self.que = que
        self.start = 0

    # def on_train_begin(self, logs=None):
    #     if self.json_path:
    #         if Path(self.json_path).exists():
    #             with open(self.json_path, 'r') as f:
    #                 self.record = json.loads(f.read())
    #             if self.start_at > 0:
    #                 for k in self.record.keys():
    #                     self.record[k] = self.record[k][:self.start_at]

    def on_epoch_end(self, epoch, logs=None):
        for k, v in logs.items():
            temp = self.record.get(k, [])
            temp.append(v)
            self.record[k] = temp
        # if self.json_path:
        #     with open(self.json_path, "w") as f:
        #         f.write(json.dumps(self.record))
        if len(self.record["loss"]) > 1:
            num = np.arange(0, len(self.record["loss"]))
            self.que.put((num, self.record))
            self.start += 1
        if self.start >= self.e_num:
            self.que.close()


def ani_loss_show(fig_path: str, que: Queue, epochs: int):
    """
    动态绘制loss值

    :param fig_path: gif图储存位置
    :param que: 进程队列
    :param epochs: epochs数量
    :return:
    """
    fig, ax = plt.subplots()
    loss_line, = plt.plot([], [], label="train_loss")
    val_loss_line, = plt.plot([], [], label="val_loss")
    acc_line, = plt.plot([], [], label="train_acc")
    val_acc_line, = plt.plot([], [], label="val_acc")
    temp_data = ([], {})

    def _init():
        plt.style.use("ggplot")
        ax.set_xlim(0, epochs)
        ax.set_ylim(0, 3)
        plt.xlabel("Epoch #")
        plt.ylabel("Loss/Accuracy")
        return loss_line, val_loss_line, acc_line, val_acc_line

    def _update(_):
        nonlocal temp_data
        try:
            num, record = que.get(block=False)
            temp_data = (num, record)
        except Empty:
            num, record = temp_data
        loss_line.set_data(num, record.get("loss", []))
        val_loss_line.set_data(num, record.get("val_loss", []))
        acc_line.set_data(num, record.get("acc", []))
        val_acc_line.set_data(num, record.get("val_acc", []))
        plt.title(f"Training Loss and Accuracy [Epoch {len(record.get('loss', []))}]")
        plt.legend()
        return loss_line, val_loss_line, acc_line, val_acc_line

    ani = FuncAnimation(fig, _update, init_func=_init)
    plt.show()
    ani.save(fig_path)

    # plt.savefig(self.fig_path)

# if __name__ == '__main__':
# import random
#
# xdata, ydata = [], []
#
#
# def init():
#     ax.set_xlim(0, 10)
#     ax.set_ylim(0, 10)
#
# def update(frame):
#     xdata.append(random.randint(0, 3))
#     ydata.append(random.randint(0, 3))
#     ln.set_data(xdata, ydata)
#
#     return ln,
#
#
# ani = FuncAnimation(fig, update, init_func=init)
# plt.show()
